#!/usr/bin/env python
"""
Program to find all taxa containing a name of interest. For example, if you want
to find all nodes with Mycobacteria in the name, this program will return the
path to those nodes on one line, and the ranks of each of those taxa on the
next.

Example:-
python findTaxa.py -in Mycobacteria Escherichia

"""

from ssummo.ssummolib import load_index
from ssummo.cmd_options import Options
from ssummo.taxonomy import find_taxa, match_ranks, walk_ranks
from sys import argv, stdout
import os.path

def print_matches(nameRanks, callback=stdout.write):
    """Traverse through the tree nameRanks, which should be a nested dictionary
    of dictionaries, printing the nodes at each level with consistent
    indentation. Results are written using the given :func:`.callback` function.
    """
    sep = os.path.sep
    taxDict = {'rank' : 'root'}
    for taxPath in nameRanks:
        node = taxDict
        taxPathList = taxPath.split( sep )
        for i, taxon in enumerate( taxPathList ):
            #taxon = taxPathList[i]
            try:
                rank = nameRanks[taxPath][i]
            except IndexError:
                node.update( { taxon : { 'rank' : nameRanks[taxPath][i-1] } } )
                callback( 'adding {0} to {1}'.format( taxon , taxPath ) )
                continue
            if taxon in node:
                node = node[taxon]
            else:
                node.update( { taxon: {'rank': rank } } )
                node = node[taxon]
    answer = stdout.write('We have {0} names matching your search criteria.'
                          .format(len(nameRanks)))
    answer = raw_input('Do you want to see them all??')
    if answer.startswith('y'):
        rankList = []
        for names, node, path in walk_ranks(taxDict):
            rank = node.pop( 'rank' )
            pathList = path.split( sep )
            depth = len( pathList )
            nRanks = len(rankList)
            if nRanks < depth:   # Update ranks
                rankList.append(set( [rank] ))
            elif nRanks == depth:
                pass
                #if rank not in rankList[depth-1]:
                #    rankList[depth-1].add(rank)
            elif rank not in rankList[depth]:
                rankList[depth].add( rank )
            else:
                pass
            nKids = len(names)
            indent = 2*depth
            if nKids == 1:
                callback( '|_{0}'.format(pathList[-1]).ljust(indent))
            elif nKids == 0:
                callback( '|-{0}'.format(pathList[-1]).ljust(indent))
            else:  # Multiple children
                callback( '|-{0}'.format(pathList[-1]).ljust(indent))
                continue
                for i in xrange(nKids-1):
                    callback( '|-{0}'.format(pathList[-1]).ljust(indent) )
                callback('|_{0}'.format(names[-1]).ljust(indent))
                #callback( '  '* depth, '|_', names[-1] )
        for i in xrange(1, len(rankList)):
            indent = 2*i
            if len(rankList[i]) == 1:
                callback('{0}{1}'.format('|-', rankList[i].pop()).ljust(indent))
            else:   # multiple
                ranks = []
                for rank in rankList[i]:
                    ranks.append(rank )
                callback('{0}{1}'.format('|-', ', '.join(ranks)).ljust(indent))

    else:
        mainPath = ['', '']
        for names, node, path in walk_ranks( taxDict ):
            if len(names) == 1:
                mainPath[0] += names[0]
                mainPath[1] += node['rank']
            else:
                break
        callback( 'The single undiverging branch we went down looked like:-' )
        callback( '\n'.join( mainPath ) )
        callback( '\nThen we had the choice of:-' )
        callback( names )
        callback( node['rank'] )

def parse_args( args ):
    options = Options()
    options.multiargs = ['-in']
    options.singleargs = []
    options.help_text = {'-in': 'Taxa name(s) of interest'}
    options.parse_args( argv[1:] )
    return options

def main(args):
    tdict = load_index()
    options = parse_args( args )
    taxaNames = find_taxa( tdict, options )
    nameRankDict = match_ranks( taxaNames )
    print_matches( nameRankDict )

if __name__ == '__main__':
    main(argv[1:])
